import os
import numpy as np
import re


def save_sub_flow(flow, output_dir, row_offset, col_offset, start_row, start_col, end_row, end_col, sub_row, sub_col,
                  tile_size=128):
    sub_flow = flow[start_row:end_row, start_col:end_col]
    filename = f"tile_({row_offset + sub_row}_{col_offset + sub_col}).npy"
    filepath = os.path.join(output_dir, filename)
    np.save(filepath, sub_flow)
    print(f"Saved {filename} to {filepath}")


def get_max_row_col(flow_files):
    max_row, max_col = 0, 0
    for flow_file in flow_files:
        match = re.search(r'tile_\((\d+)_(\d+),(\d+)_(\d+)\)\.npy', flow_file)
        if match:
            row, col, _, _ = map(int, match.groups())
            if row > max_row:
                max_row = row
            if col > max_col:
                max_col = col
    return max_row, max_col


def split_and_save_flow(flow_dir, output_dir, tile_size=128):
    os.makedirs(output_dir, exist_ok=True)

    flow_files = [f for f in os.listdir(flow_dir) if f.endswith('.npy')]
    max_row, max_col = get_max_row_col(flow_files)

    for flow_file in flow_files:
        flow_path = os.path.join(flow_dir, flow_file)
        flow = np.load(flow_path)

        h, w, _ = flow.shape

        match = re.search(r'tile_\((\d+)_(\d+),(\d+)_(\d+)\)\.npy', flow_file)

        if not match:
            raise ValueError(f"文件名格式不正确: {flow_file}")

        top_left_row, top_left_col, _, _ = map(int, match.groups())

        # 计算行列偏移
        if top_left_row == 0:
            row_offset = 0
        else:
            row_offset = top_left_row * 2

        if top_left_col == 0:
            col_offset = 0
        else:
            col_offset = top_left_col * 2

        range_i = [1, 2]
        range_j = [1, 2]
        if top_left_row == 0:  # 第一行
            range_i.append(0)

        if top_left_row == max_row:  # 最后一行
            range_i.append(3)

        if top_left_col == 0:  # 第一列
            range_j.append(0)

        if top_left_col == max_col:  # 最后一列
            range_j.append(3)

        # 划分光流场为128x128的小块
        for i in range(4):
            for j in range(4):
                start_row, start_col = i * tile_size, j * tile_size
                end_row, end_col = start_row + tile_size, start_col + tile_size

                # 中间部分
                if i in range_i and j in range_j:
                    save_sub_flow(flow, output_dir, row_offset, col_offset, start_row, start_col, end_row, end_col, i,
                                  j)


# 示例用法
flow_dir = r'E:\wafer52\256nm_flow_combine_512'  # 替换为实际路径
output_dir = r'E:\wafer52\256nm_flow_split_64'  # 替换为实际路径

split_and_save_flow(flow_dir, output_dir, tile_size=64)

print(f"Processed flow fields saved to {output_dir}")
